package com.by.rpc.netty;

import com.by.log.Logger;
import com.by.log.LoggerFactory;
import com.by.log.TraceContext;
import com.by.rpc.dto.RpcResponse;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;

import java.util.Map;
import java.util.concurrent.CompletableFuture;

/**
 * Netty客户端处理器，用于处理RPC响应
 */
public class RpcClientHandler extends ChannelInboundHandlerAdapter {
    private static final Logger log = LoggerFactory.getLogger(RpcClientHandler.class);
    private final Map<String, CompletableFuture<RpcResponse>> requestMap;

    public RpcClientHandler(Map<String, CompletableFuture<RpcResponse>> requestMap) {
        this.requestMap = requestMap;
        log.info("RpcClientHandler initialized");
    }

    // 接收的msg直接是RpcResponse对象（Kryo已反序列化）
    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
        RpcResponse rpcResponse = (RpcResponse) msg;  // 直接强转
        log.debug("Received RPC response: {}", rpcResponse);
        
        // 设置追踪上下文
        if (rpcResponse.getTraceId() != null) {
            TraceContext traceContext = TraceContext.getCurrentContext();
            traceContext.setTraceId(rpcResponse.getTraceId());
            traceContext.setSpanId(rpcResponse.getSpanId());
            traceContext.setParentSpanId(rpcResponse.getParentSpanId());
        }

        String requestId = rpcResponse.getRequestId();
        CompletableFuture<RpcResponse> responseFuture = requestMap.remove(requestId);
        if (responseFuture != null) {
            if (rpcResponse.getStatus() == RpcResponse.ResponseStatus.SUCCESS) {
                responseFuture.complete(rpcResponse);
            } else {
                responseFuture.completeExceptionally(new RuntimeException(rpcResponse.getErrorMsg()));
            }
        } else {
            log.warn("Received response for unknown request: {}", requestId);
        }
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        log.error("RPC client handler exception: {}", cause.getMessage(), cause);
        // 异常时完成所有未完成的请求
        requestMap.values().forEach(future -> future.completeExceptionally(cause));
        requestMap.clear();
        ctx.close();
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        log.warn("RPC client channel inactive");
        super.channelInactive(ctx);
    }

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception {
        log.info("RPC client channel active");
        super.channelActive(ctx);
    }
}